package 困难.并查集;

import java.util.Arrays;
import java.util.PriorityQueue;

/**
 * 给你一个大小为 m x n 的整数矩阵 grid 和一个大小为 k 的数组 queries 。
 * 找出一个大小为 k 的数组 answer ，且满足对于每个整数 queres[i] ，你
 * 从矩阵 左上角 单元格开始，重复以下过程：
 * 如果 queries[i] 严格 大于你当前所处位置单元格，如果该单元格是第一次访
 * 问，则获得 1 分，并且你可以移动到所有 4 个方向（上、下、左、右）上任一 相邻 单元格。
 * 否则，你不能获得任何分，并且结束这一过程。
 * 在过程结束后，answer[i] 是你可以获得的最大分数。注意，对于每个查询，你可以访问同一个单元格 多次 。
 * 返回结果数组 answer 。
 * <p>
 * 来源：力扣（LeetCode）
 * 链接：https://leetcode.cn/problems/maximum-number-of-points-from-grid-queries
 */
public class 矩阵查询可获得的最大分数_2503 {

    public static void main(String[] args) {

        System.out.println(Arrays.toString(maxPoints2(new int[][]{{1, 2, 3}, {2, 5, 7}, {3, 5, 1}},
                new int[]{5, 6, 2})));

    }

    /**
     * 并查集+离线询问
     * 1，将格子看成连通图，每次查询与 [0,0] 连通所有连通的节点数量即为该查询的答案
     * 2，将查询排序，这样与 [0,0] 一定是逐渐扩张的
     * 3，两个点的连接条件为 query[i] > Math.max(a,b)
     * 4，并查集用一个 size[i] 表示以 i 为根节点连通的所有节点数量
     * 5，离线询问，将查询从小到大排序，保存其对应的索引，将格子里的数字也从小
     * 到大排序，保存对应的索引位置，使用双指针遍历排序后的格子，即优先连通
     * 那些数字较小的格子，可能会连通到 [0,0]对应的块，也可能不连通 [0,0]
     */
    public static int[] maxPoints(int[][] grid, int[] queries) {
        int[] ans = new int[queries.length];
        int[][] q = new int[queries.length][2];
        for (int i = 0; i < queries.length; i++) {
            q[i] = new int[]{i, queries[i]};
        }
        Arrays.sort(q, (a, b) -> a[1] - b[1]);

        int m = grid.length, n = grid[0].length;
        int[][] g = new int[m * n][3];
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                g[i * n + j] = new int[]{grid[i][j], i, j};
            }
        }
        Arrays.sort(g, (a, b) -> a[0] - b[0]);

        DSU dsu = new DSU(m * n);
        int[][] direction = new int[][]{{0, 1}, {0, -1}, {1, 0}, {-1, 0}};

        int gIdx = 0;
        for (int i = 0; i < q.length; i++) {
            for (; gIdx < m * n && g[gIdx][0] < q[i][1]; gIdx++) {
                for (int[] direct : direction) {
                    int nextI = g[gIdx][1] + direct[0];
                    int nextJ = g[gIdx][2] + direct[1];
                    if (nextI >= 0 && nextI < m && nextJ >= 0 && nextJ < n &&
                            grid[nextI][nextJ] < q[i][1]) {
                        dsu.union(g[gIdx][1] * n + g[gIdx][2], nextI * n + nextJ);
                    }
                }
            }

            if (grid[0][0] < q[i][1]) {
                ans[q[i][0]] = dsu.size[dsu.find(0)];
            }
        }
        return ans;
    }

    static class DSU {

        private int[] pre;

        private int[] size;

        public DSU(int n) {
            pre = new int[n];
            size = new int[n];
            for (int i = 0; i < n; i++) {
                pre[i] = i;
            }
            Arrays.fill(size, 1);
        }

        public int find(int x) {
            if (pre[x] == x) {
                return x;
            }
            return pre[x] = find(pre[x]);
        }

        public void union(int x, int y) {
            x = find(x);
            y = find(y);

            if (x == y) return;

            pre[x] = y;
            size[y] += size[x];
        }
    }

    /**
     * 优先队列（最小堆）
     * 1，查询排序
     * 2，队列中 最小值排在队首，如果堆顶元素小于当前查询，记录从队列中取出来的累计个数，即为当前查询的结果
     * 3，1 <= grid[i][j] ，遍历过的 grid[i][j] 置为 0，防止重复访问
     */
    public static int[] maxPoints2(int[][] grid, int[] queries) {
        int[][] direction = new int[][]{{0, 1}, {0, -1}, {1, 0}, {-1, 0}};

        int[][] q = new int[queries.length][2];
        for (int i = 0; i < q.length; i++) {
            q[i] = new int[]{i, queries[i]};
        }
        Arrays.sort(q, (a, b) -> a[1] - b[1]);

        // 优先队列里有值的缓存，加入队列后可以将 grid 里的值置为 0
        PriorityQueue<int[]> priorityQueue = new PriorityQueue<>((a, b) -> a[0] - b[0]);
        priorityQueue.add(new int[]{grid[0][0], 0, 0});
        grid[0][0] = 0;
        int[] ans = new int[queries.length];
        int count = 0;

        int m = grid.length, n = grid[0].length;
        for (int i = 0; i < q.length; i++) {
            while (!priorityQueue.isEmpty() && priorityQueue.peek()[0] < q[i][1]) {
                int[] curNode = priorityQueue.poll();
                count++;

                for (int[] direct : direction) {
                    int nextRow = curNode[1] + direct[0];
                    int nextColumn = curNode[2] + direct[1];
                    if (nextRow >= 0 && nextRow < m && nextColumn >= 0 && nextColumn < n &&
                            grid[nextRow][nextColumn] != 0) {
                        priorityQueue.add(new int[]{grid[nextRow][nextColumn], nextRow, nextColumn});
                        grid[nextRow][nextColumn] = 0;
                    }
                }
            }
            ans[q[i][0]] = count;
        }
        return ans;
    }
}
